package com.example.studyllm.ollama;

import cn.hutool.json.JSONUtil;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.loader.UrlDocumentLoader;
import dev.langchain4j.data.document.parser.apache.tika.ApacheTikaDocumentParser;
import dev.langchain4j.data.document.splitter.DocumentByLineSplitter;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.ollama.OllamaChatModel;
import dev.langchain4j.model.ollama.OllamaEmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingSearchResult;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.filter.Filter;
import dev.langchain4j.store.embedding.pgvector.PgVectorEmbeddingStore;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;

import static dev.langchain4j.store.embedding.filter.MetadataFilterBuilder.metadataKey;

/**
 * @author wk
 */
@RestController
@RequestMapping("/ollama")
public class EmbeddingController {

    @GetMapping(value="/embed")
    public String embed() {
        Document document;
        document = FileSystemDocumentLoader.loadDocument("E:\\新建文本文档.txt", new ApacheTikaDocumentParser());
        document.metadata().put("fileName", "c.md");
        DocumentSplitter splitter = new DocumentByLineSplitter(100,0);
        List<TextSegment> segments = splitter.split(document);


        EmbeddingModel embeddingModel = buildEmbedding();
        EmbeddingStore<TextSegment> embeddingStore = buildEmbeddingStore();
        List<Embedding> embeddings = embeddingModel.embedAll(segments).content();
        List<String> ids = embeddingStore.addAll(embeddings, segments);
        // 正则表达式匹配换行符
        return JSONUtil.toJsonStr(ids);
    }
    @GetMapping(value="/search")
    public String search() {

        EmbeddingModel embeddingModel = buildEmbedding();
        EmbeddingStore<TextSegment> embeddingStore = buildEmbeddingStore();
        Embedding queryEmbedding = embeddingModel.embed("MySQL创建语句").content();
        Filter filter = metadataKey("fileName").isEqualTo("c.md");
        EmbeddingSearchResult<TextSegment> list = embeddingStore.search(EmbeddingSearchRequest
                .builder()
                .queryEmbedding(queryEmbedding)
                .maxResults(5)
                .filter(filter)
                .build());

        List<Map<String, Object>> result = new ArrayList<>();
        list.matches().forEach(i -> {
            TextSegment embedded = i.embedded();
            Map<String, Object> map = embedded.metadata().toMap();
            map.put("text", embedded.text());
            result.add(map);
        });

        String promot = """
                查询MySQL创建语句，
                以下是文本内容，请根据内容提取问题的结果：
                """ + JSONUtil.toJsonStr(result);
        ChatLanguageModel model = buildModel();

        return  model.chat(promot);
    }


    private ChatLanguageModel buildModel(){
        return OllamaChatModel.builder()
                .baseUrl("http://47.109.192.172:11434")
                .modelName("qwen2:7b")
                .temperature(0.1)
                .build();
    }

    public EmbeddingModel buildEmbedding() {
            return OllamaEmbeddingModel
                    .builder()
                    .baseUrl("http://47.109.192.172:11434")
                    .modelName("nomic-embed-text")
                    .logRequests(true)
                    .logResponses(true)
                    .build();
    }

    private EmbeddingStore buildEmbeddingStore() {
        PgVectorEmbeddingStore store = PgVectorEmbeddingStore.builder()
                .host("47.109.192.172")
                .port(5432)
                .database("langchat")
                .dimension(768)
                .user("root")
                .password("root")
                .table("testEmb")
                .indexListSize(1)
                .useIndex(true)
                .createTable(true)
                .dropTableFirst(false)
                .build();
        return store;
    }
}
